package bus
import chisel3._
import chisel3.util._
import utils.HoldUnless

class AXI4Crossbar1toN(addressSpace: List[(Long, Long)]) extends Module{
  val io = IO(new Bundle() {
    val in = Flipped(new AXI4)
    val out = Vec(addressSpace.length, new AXI4)
  })

  val s_idle :: s_resp :: s_error :: Nil = Enum(3)
  val state = RegInit(s_idle)

  // select the output channel according to the address
  val addr = io.in.ar.bits.addr
  val outSelVec = VecInit(addressSpace.map(range => (addr >= range._1.U && addr < (range._1 + range._2).U)))
  val outSelIdx = PriorityEncoder(outSelVec)
  val outSel = io.out(outSelIdx)
  val outSelIdxResp = RegEnable(outSelIdx, outSel.ar.fire() && (state === s_idle))
  val outSelResp = io.out(outSelIdxResp)
  val reqInvalidAddr = io.in.ar.valid && !outSelVec.asUInt.orR

  // bind out.req channel
  (io.out zip outSelVec).map { case (o, v) => {
    o.ar.bits := io.in.ar.bits
    o.ar.valid := v && (io.in.ar.valid && (state === s_idle))
    o.r.ready := v
  }}
  switch (state) {
    is (s_idle) {
      when (outSel.ar.fire()) { state := s_resp }
      when (reqInvalidAddr) { state := s_error }
    }
    is (s_resp) { when (outSelResp.r.fire()) { state := s_idle } }
    is (s_error) { when(io.in.r.fire()){ state := s_idle } }
  }

  io.in.r.valid := outSelResp.r.fire() || state === s_error
  io.in.r.bits <> outSelResp.r.bits
  outSelResp.r.ready := io.in.r.ready
  io.in.ar.ready := outSel.ar.ready || reqInvalidAddr

  /**
   * write
   */
  val w_idle :: w_write :: w_resp :: w_error :: Nil = Enum(4)
  val w_state = RegInit(w_idle)
  // select the output channel according to the address
  val waddr = io.in.aw.bits.addr
  val woutSelVec = VecInit(addressSpace.map(range => (waddr >= range._1.U && waddr < (range._1 + range._2).U)))
  val woutSelIdx = PriorityEncoder(woutSelVec)
  val woutSel = io.out(woutSelIdx)
  val woutSelIdxResp = HoldUnless(woutSelIdx, woutSel.aw.fire() && (w_state === w_idle))
  val woutSelResp = io.out(woutSelIdxResp)
  val wreqInvalidAddr = io.in.aw.valid && !woutSelVec.asUInt.orR

  // bind out.req channel
  (io.out zip woutSelVec).map { case (o, v) => {
    o.aw.bits := io.in.aw.bits
    o.aw.valid := v && (io.in.aw.valid && (w_state === w_idle))
    o.b.ready := v
  }}

  io.in.w.ready := woutSelResp.w.ready
  io.out.zipWithIndex.map{ case (o, i) => {
    o.w.valid := io.in.w.valid && (woutSelIdxResp === i.U)
    o.w.bits := io.in.w.bits
  }}

  switch (w_state) {
    is (w_idle) {
      when (woutSel.aw.fire()) { w_state := w_write }
      when (woutSelResp.w.fire() && woutSelResp.w.bits.last) {w_state := w_resp}
      when (wreqInvalidAddr) { w_state := w_error }
    }
    is(w_write) { when(woutSelResp.w.fire() && woutSelResp.w.bits.last) { w_state := w_resp }  }
    is (w_resp) { when (woutSelResp.b.fire()) { w_state := w_idle } }
    is (w_error) { when(io.in.b.fire()){ w_state := w_idle } }
  }
  io.in.b.valid := woutSelResp.b.fire() || w_state === w_error
  io.in.b.bits <> woutSelResp.b.bits

  woutSelResp.b.ready := io.in.b.ready
  io.in.aw.ready := woutSel.aw.ready || wreqInvalidAddr

}

class AXI4CrossbarNto1(n: Int) extends Module{
  val io = IO(new Bundle() {
    val in = Vec(n, Flipped(new AXI4))
    val out = new AXI4
  })

  val r_arb = Module(new Arbiter(new AXI4BundleA(), n))
  val w_arb = Module(new Arbiter(new AXI4BundleA(), n))

  val r_chosen = Wire(UInt(n.W))
  val w_chosen = Wire(UInt(n.W))
  val r_port = RegInit(0.U(log2Up(n).W))
  val w_port = RegInit(0.U(log2Up(n).W))
  for(i <- 0 until n){
    r_arb.io.in(i) <> io.in(i).ar
    w_arb.io.in(i) <> io.in(i).aw
  }

  r_chosen := r_arb.io.chosen
  w_chosen := w_arb.io.chosen
  val thisReadReq = r_arb.io.out
  val thisWriteReq = w_arb.io.out

  val rIdle :: memReadResp :: Nil = Enum(2)
  val r_state = RegInit(rIdle)

  /**
   * read events
   */
  io.out.ar.valid := thisReadReq.valid && (r_state === rIdle)
  io.out.ar.bits := thisReadReq.bits
  thisReadReq.ready := io.out.ar.ready && (r_state === rIdle)

  io.in.map(_.r.bits := io.out.r.bits)
  io.in.map(_.r.valid := false.B)

  io.in(r_port).r.valid := io.out.r.valid
  io.out.r.ready := io.in(r_port).r.ready

  switch(r_state) {
    is(rIdle) {
      when(thisReadReq.fire()) {
        r_port := r_chosen
        r_state := memReadResp
      }
    }
    is(memReadResp) {
      when(io.in(r_port).r.fire() && io.in(r_port).r.bits.last){
        r_state := rIdle
      }
    }

  }

  /**
   * write events
   */
  val wIdle :: memWrite :: memWriteB :: Nil = Enum(3)
  val w_state = RegInit(wIdle)

  // write addr
  io.out.aw.valid := thisWriteReq.valid && (w_state === wIdle)
  io.out.aw.bits := thisWriteReq.bits
  thisWriteReq.ready := io.out.aw.ready && (w_state === wIdle)

  w_port := HoldUnless(w_chosen, thisWriteReq.fire())
  // write data
  io.in.map(_.w.ready := false.B)
  io.out.w.valid := io.in(w_port).w.valid
  io.out.w.bits := io.in(w_port).w.bits
  io.in(w_port).w.ready := io.out.w.ready && ((w_state === memWrite) || (w_state === wIdle))

  // write response
  io.in.map(_.b.bits := io.out.b.bits)
  io.in.map(_.b.valid := false.B)

  io.in(w_port).b.valid := io.out.b.valid
  io.out.b.ready := io.in(w_port).b.ready

  switch(w_state) {
    is(wIdle) {
      when(thisWriteReq.fire()) {
        when(io.out.w.fire() && io.out.w.bits.last){
          w_state := memWriteB
        }.otherwise{ w_state := memWrite }

      }
    }
    is(memWrite) {
      when(io.out.w.fire() && io.out.w.bits.last){
        w_state := memWriteB
      }
    }
    is(memWriteB) {
      when(io.in(w_port).b.fire()){
        w_state := wIdle
      }
    }
  }
}
